#pragma once

#include <common/HeaderSpace.h>
#include <common/Port.h>
#include <common/Rule.h>

#include <map>
#include <memory>
#include <string>
#include <vector>

namespace flashpp {

// class Device;
// class Rule;

class Index {
 public:
  Index() {
    root_ = std::make_shared<TrieNode>(0);
  }

  struct TrieNode {
    TrieNode(int level) : level(level) {
      left = right = nullptr;
      rule = nullptr;
    }
    int level;
    std::shared_ptr<Rule> rule;
    std::shared_ptr<TrieNode> left, right;
  };

  void
  insert(std::shared_ptr<Rule> rule) {
    insert(rule, root_);
  }

  void
  insert(std::shared_ptr<Rule> rule, std::shared_ptr<TrieNode> node) {
    if (node->level == rule->prefixLength) {
      // if (node->rule != nullptr) {
      //   std::cout << "exist" << std::endl;
      //   std::cout << node->rule->addr << ":" << node->rule->prefixLength << std::endl;
      //   std::cout << rule->addr << ":" << rule->prefixLength << std::endl;
      // }
      node->rule = rule;
      return;
    }
    int level = node->level;
    if (rule->bits[31 - level] == 0) {
      if (node->left == nullptr) {
        node->left = std::make_shared<TrieNode>(level + 1);
      }
      insert(rule, node->left);
    } else {
      if (node->right == nullptr) {
        node->right = std::make_shared<TrieNode>(level + 1);
      }
      insert(rule, node->right);
    }
  }

  void
  getOverlapSet(
      std::shared_ptr<Rule> rule, std::vector<std::shared_ptr<Rule>>& out) {
    getOverlapSet(rule, root_, out);
  }

  void
  getOverlapSet(
      std::shared_ptr<Rule> rule,
      std::shared_ptr<TrieNode> node,
      std::vector<std::shared_ptr<Rule>>& out) {
    if (node == nullptr) {
      return;
    }
    if (node->level > rule->prefixLength &&
        node->rule != nullptr) { // we only return rules with higher priorities
      out.push_back(node->rule);
    }
    if (node->level < rule->prefixLength) {
      if (rule->bits[31 - node->level] == 0) {
        getOverlapSet(rule, node->left, out);
      } else {
        getOverlapSet(rule, node->right, out);
      }
    } else {
      getOverlapSet(rule, node->left, out);
      getOverlapSet(rule, node->right, out);
    }
  }

 private:
  std::shared_ptr<TrieNode> root_;
};

class Device {
 public:
  Device(std::string name) : name_(name){};

  /**
   * Get or create a port if port does not exists and return the port reference
   */
  Port& getCreatePort(std::string name);

  std::string&
  name() {
    return name_;
  }

  auto&
  eportToHs() {
    return eportToHs_;
  }

  auto&
  eportToHs(Port* eport) {
    return eportToHs_.at(eport);
  }

  auto&
  eportToAps() {
    return eportToAps_;
  }

  // void addEdge(std::string eport, std::string dstdev) {
  //   eportToEdge_.emplace(std::piecewise_construct,
  //   std::forward_as_tuple(eport), std::forward_as_tuple(dstdev));
  // }

  void
  addNeighbor(std::string pn, std::string neighbor) {
    portToNeighbor_[pn] = neighbor;
  }

  std::string
  getNeighbor(std::string pn) {
    if (portToNeighbor_.find(pn) == portToNeighbor_.end()) {
      return "";
    }
    return portToNeighbor_[pn];
  }

  std::map<std::string, Port>&
  nameToPort() {
    return nameToPort_;
  }

  auto&
  rules() {
    return rules_;
  }

  void
  addRule(std::shared_ptr<Rule> rule) {
    rules_.push_back(rule);
    index_.insert(rule);
  }

  void
  getOverlapSet(
      std::shared_ptr<Rule> rule, std::vector<std::shared_ptr<Rule>>& out) {
    index_.getOverlapSet(rule, out);
  }

 private:
  std::string name_;
  // std::map<int, Rule, std::greater<int>> rules_;
  // std::vector<Rule> rules_;
  std::vector<std::shared_ptr<Rule>> rules_;

  std::map<Port*, HeaderSpace> eportToHs_;

  std::map<std::string, std::vector<HeaderSpace>> eportToAps_;

  std::map<std::string, Port> nameToPort_;

  std::map<std::string, std::string> portToNeighbor_;

  Index index_;
};

} // namespace flashpp